{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Title: #Minimum Sum of Squared Difference"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Difficulty: #Medium"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Category Title: #Algorithms"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Tag Slug: #array #math #sorting #heap-priority-queue"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Name Translated: #数组 #数学 #排序 #堆（优先队列）"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Solution Name: minSumSquareDiff"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Translated Title: #最小差值平方和"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Translated Content:\n",
    "<p>给你两个下标从 <strong>0</strong>&nbsp;开始的整数数组&nbsp;<code>nums1</code> 和&nbsp;<code>nums2</code>&nbsp;，长度为&nbsp;<code>n</code>&nbsp;。</p>\n",
    "\n",
    "<p>数组&nbsp;<code>nums1</code> 和&nbsp;<code>nums2</code>&nbsp;的 <strong>差值平方和</strong>&nbsp;定义为所有满足&nbsp;<code>0 &lt;= i &lt; n</code>&nbsp;的&nbsp;<code>(nums1[i] - nums2[i])<sup>2</sup></code>&nbsp;之和。</p>\n",
    "\n",
    "<p>同时给你两个正整数&nbsp;<code>k1</code> 和&nbsp;<code>k2</code>&nbsp;。你可以将&nbsp;<code>nums1</code>&nbsp;中的任意元素&nbsp;<code>+1</code> 或者&nbsp;<code>-1</code>&nbsp;至多&nbsp;<code>k1</code>&nbsp;次。类似的，你可以将&nbsp;<code>nums2</code>&nbsp;中的任意元素&nbsp;<code>+1</code> 或者&nbsp;<code>-1</code>&nbsp;至多&nbsp;<code>k2</code>&nbsp;次。</p>\n",
    "\n",
    "<p>请你返回修改数组<em>&nbsp;</em><code>nums1</code><em>&nbsp;</em>至多<em>&nbsp;</em><code>k1</code>&nbsp;次且修改数组<em>&nbsp;</em><code>nums2</code>&nbsp;至多 <code>k2</code><em>&nbsp;</em>次后的最小&nbsp;<strong>差值平方和</strong>&nbsp;。</p>\n",
    "\n",
    "<p><strong>注意：</strong>你可以将数组中的元素变成&nbsp;<strong>负</strong>&nbsp;整数。</p>\n",
    "\n",
    "<p>&nbsp;</p>\n",
    "\n",
    "<p><strong>示例 1：</strong></p>\n",
    "\n",
    "<pre><b>输入：</b>nums1 = [1,2,3,4], nums2 = [2,10,20,19], k1 = 0, k2 = 0\n",
    "<b>输出：</b>579\n",
    "<b>解释：</b>nums1 和 nums2 中的元素不能修改，因为 k1 = 0 和 k2 = 0 。\n",
    "差值平方和为：(1 - 2)<sup>2 </sup>+ (2 - 10)<sup>2 </sup>+ (3 - 20)<sup>2 </sup>+ (4 - 19)<sup>2</sup>&nbsp;= 579 。\n",
    "</pre>\n",
    "\n",
    "<p><strong>示例 2：</strong></p>\n",
    "\n",
    "<pre><b>输入：</b>nums1 = [1,4,10,12], nums2 = [5,8,6,9], k1 = 1, k2 = 1\n",
    "<b>输出：</b>43\n",
    "<b>解释：</b>一种得到最小差值平方和的方式为：\n",
    "- 将 nums1[0] 增加一次。\n",
    "- 将 nums2[2] 增加一次。\n",
    "最小差值平方和为：\n",
    "(2 - 5)<sup>2 </sup>+ (4 - 8)<sup>2 </sup>+ (10 - 7)<sup>2 </sup>+ (12 - 9)<sup>2</sup>&nbsp;= 43 。\n",
    "注意，也有其他方式可以得到最小差值平方和，但没有得到比 43 更小答案的方案。</pre>\n",
    "\n",
    "<p>&nbsp;</p>\n",
    "\n",
    "<p><strong>提示：</strong></p>\n",
    "\n",
    "<ul>\n",
    "\t<li><code>n == nums1.length == nums2.length</code></li>\n",
    "\t<li><code>1 &lt;= n &lt;= 10<sup>5</sup></code></li>\n",
    "\t<li><code>0 &lt;= nums1[i], nums2[i] &lt;= 10<sup>5</sup></code></li>\n",
    "\t<li><code>0 &lt;= k1, k2 &lt;= 10<sup>9</sup></code></li>\n",
    "</ul>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Description: [minimum-sum-of-squared-difference](https://leetcode.cn/problems/minimum-sum-of-squared-difference/description/)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Solutions: [minimum-sum-of-squared-difference](https://leetcode.cn/problems/minimum-sum-of-squared-difference/solutions/)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_cases = ['[1,2,3,4]\\n[2,10,20,19]\\n0\\n0', '[1,4,10,12]\\n[5,8,6,9]\\n1\\n1']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minSumSquareDiff(self, nums1: List[int], nums2: List[int], k1: int, k2: int) -> int:\n",
    "        cnt = [0] * 100001\n",
    "        for nums1, nums2 in zip(nums1, nums2):\n",
    "            cnt[abs(nums1 - nums2)] += 1\n",
    "        k = k1 + k2\n",
    "        for i in range(len(cnt) - 1, 0, -1):\n",
    "            change = min(k, cnt[i])\n",
    "            cnt[i - 1] += change\n",
    "            k -= change\n",
    "            cnt[i] -= change\n",
    "        res = 0\n",
    "        for i in range(len(cnt)):\n",
    "            if cnt[i] > 0:\n",
    "                res += cnt[i] * i * i\n",
    "        return res\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minSumSquareDiff(self, a: List[int], nums2: List[int], k1: int, k2: int) -> int:\n",
    "        ans, k = 0, k1 + k2\n",
    "        for i in range(len(a)):\n",
    "            a[i] = abs(a[i] - nums2[i])\n",
    "            ans += a[i] * a[i]\n",
    "        if sum(a) <= k:\n",
    "            return 0  # 所有 a[i] 均可为 0\n",
    "        a.sort(reverse=True)\n",
    "        a.append(0)  # 哨兵\n",
    "        for i, v in enumerate(a):\n",
    "            ans -= v * v\n",
    "            j = i + 1\n",
    "            c = j * (v - a[j])\n",
    "            if c < k:\n",
    "                k -= c\n",
    "                continue\n",
    "            v -= k // j\n",
    "            return ans + k % j * (v - 1) * (v - 1) + (j - k % j) * v * v"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minSumSquareDiff(self, nums1: List[int], nums2: List[int], k1: int, k2: int) -> int:\n",
    "        count = [0] * (100001)\n",
    "        for i in range(len(nums1)):\n",
    "            count[abs(nums1[i] - nums2[i])] += 1\n",
    "        k = k1 + k2\n",
    "        for i in range(len(count) - 1, 0, -1):\n",
    "            change = min(k, count[i])\n",
    "            count[i - 1] += change\n",
    "            k -= change\n",
    "            count[i] -= change\n",
    "        res = 0\n",
    "        for i in range(len(count)):\n",
    "            if count[i] != 0:\n",
    "                res += count[i] * i * i \n",
    "        return res\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minSumSquareDiff(self, nums1: List[int], nums2: List[int], k1: int, k2: int) -> int:\n",
    "        k=k1+k2\n",
    "        n=len(nums1)\n",
    "        dif=[0]*n\n",
    "        for i in range(n):\n",
    "            dif[i]=abs(nums1[i]-nums2[i])\n",
    "        dif.sort(reverse=True)\n",
    "        def check(t):\n",
    "            ans=0\n",
    "            for i in range(n):\n",
    "                if dif[i]>t:\n",
    "                    ans+=dif[i]-t\n",
    "                else:\n",
    "                    break\n",
    "            if ans<=k:\n",
    "                return True\n",
    "            else:\n",
    "                return False\n",
    "        l=0\n",
    "        r=max(dif)\n",
    "        #print(dif)\n",
    "        while l+1<r:\n",
    "            mid=(l+r)//2\n",
    "            if check(mid):\n",
    "                r=mid\n",
    "            else:\n",
    "                l=mid\n",
    "        \n",
    "        if check(l):\n",
    "            y=l\n",
    "        else:\n",
    "            y=r\n",
    "        #print(dif)\n",
    "        #print(y)\n",
    "        ans=0\n",
    "        if y==0:\n",
    "            return 0\n",
    "\n",
    "        for i in range(n):\n",
    "            if dif[i]>y:\n",
    "                ans+=dif[i]-y\n",
    "                dif[i]=y\n",
    "        k =k-ans\n",
    "      #  print(k)\n",
    "        dif.sort(reverse=True)\n",
    "        for i in range(k):\n",
    "            dif[i]=dif[i]-1\n",
    "        res=0\n",
    "        for i in range(n):\n",
    "            res+=dif[i]**2\n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minSumSquareDiff(self, nums1: List[int], nums2: List[int], k1: int, k2: int) -> int:\n",
    "        n = len(nums1)\n",
    "        for i in range(0, n ):\n",
    "            nums1[i] = abs(nums1[i] - nums2[i])\n",
    "        nums1.sort()\n",
    "        maxx = nums1[n - 1]\n",
    "        num = 1\n",
    "        flag = n - 2\n",
    "        last = k1 + k2\n",
    "        while flag >=0:\n",
    "            need = num * (maxx - nums1[flag])\n",
    "            if need >= last:\n",
    "                cut = int(last / num)\n",
    "                maxx -= cut\n",
    "                last = last - cut * num\n",
    "                ans = pow(maxx - 1, 2) * last + (num - last) * pow(maxx, 2)\n",
    "                for j in range(0, flag + 1):\n",
    "                    ans += pow(nums1[j],2)\n",
    "                return ans\n",
    "            else:\n",
    "                last -= num * (maxx - nums1[flag])\n",
    "                maxx = nums1[flag]\n",
    "                num += 1\n",
    "                flag -= 1\n",
    "        if maxx * num <= last:\n",
    "            return 0\n",
    "        cut = int(last / num)\n",
    "        maxx -= cut\n",
    "        last = last - cut * num\n",
    "        ans = pow(maxx - 1, 2) * last + (num - last) * pow(maxx, 2)\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minSumSquareDiff(self, nums1: List[int], nums2: List[int], k1: int, k2: int) -> int:\n",
    "        diff = []\n",
    "        n = len(nums1)\n",
    "        res = 0\n",
    "        for i in range(n):\n",
    "            diff.append(abs(nums1[i]-nums2[i]))\n",
    "            res += diff[i] * diff[i]\n",
    "        k = k1 + k2\n",
    "        if sum(diff) <= k:\n",
    "            return 0\n",
    "        diff.sort(reverse=True)\n",
    "        diff.append(0)\n",
    "        for i, v in enumerate(diff):\n",
    "            res -= v * v\n",
    "            c = i + 1\n",
    "            cnt = c * (v - diff[c])\n",
    "            if cnt < k:\n",
    "                k -= cnt\n",
    "                continue\n",
    "            v -= k // c\n",
    "            res += (c - k % c) * v * v\n",
    "            res +=  k % c * (v - 1) * (v - 1)\n",
    "            return res  \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from sortedcontainers import SortedList\n",
    "class Solution:\n",
    "    def minSumSquareDiff(self, nums1: List[int], nums2: List[int], k1: int, k2: int) -> int:\n",
    "        n = len(nums1)\n",
    "        delta = [abs(x1 - x2) for x1, x2 in zip(nums1, nums2)]\n",
    "        # print(delta)\n",
    "        \n",
    "        def check(num):\n",
    "            cnt = 0\n",
    "            for i in range(n):\n",
    "                cnt += max(0, delta[i] - num)\n",
    "            return cnt <= k1 + k2\n",
    "        \n",
    "        i = 0\n",
    "        j = max(delta)\n",
    "        while i < j:\n",
    "            mid = (i + j) >> 1\n",
    "            if check(mid):\n",
    "                j = mid\n",
    "            else:\n",
    "                i = mid + 1\n",
    "        \n",
    "        # print(j)\n",
    "        ans = 0\n",
    "        cnt = 0\n",
    "        for i in range(n):\n",
    "            cnt += max(0, delta[i] - j)\n",
    "            ans += min(delta[i], j) ** 2 \n",
    "        return ans - (k1 + k2 - cnt) * (2 * j - 1) if j > 0 else 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minSumSquareDiff(self, nums1: List[int], nums2: List[int], k1: int, k2: int) -> int:\n",
    "        count = [0] * (100001)\n",
    "        for i in range(len(nums1)):\n",
    "            count[abs(nums1[i] - nums2[i])] += 1\n",
    "        k = k1 + k2\n",
    "        for i in range(len(count) - 1, 0, -1):\n",
    "            change = min(k, count[i])\n",
    "            count[i - 1] += change\n",
    "            k -= change\n",
    "            count[i] -= change\n",
    "        res = 0\n",
    "        for i in range(len(count)):\n",
    "            if count[i] != 0:\n",
    "                res += count[i] * i * i \n",
    "        return res\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minSumSquareDiff(self, nums1: List[int], nums2: List[int], k1: int, k2: int) -> int:\n",
    "        cnt = [0] * 100001\n",
    "        for nums1, nums2 in zip(nums1, nums2):\n",
    "            cnt[abs(nums1 - nums2)] += 1\n",
    "        k = k1 + k2\n",
    "        for i in range(len(cnt) - 1, 0, -1):\n",
    "            change = min(k, cnt[i])\n",
    "            cnt[i - 1] += change\n",
    "            k -= change\n",
    "            cnt[i] -= change\n",
    "        res = 0\n",
    "        for i in range(len(cnt)):\n",
    "            if cnt[i] > 0:\n",
    "                res += cnt[i] * i * i\n",
    "        return res\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minSumSquareDiff(self, nums1: List[int], nums2: List[int], k1: int, k2: int) -> int:\n",
    "        v = [abs(x - y) for x, y in zip(nums1, nums2)]\n",
    "        l, r = -1, max(v) + 1\n",
    "        while l + 1 < r:\n",
    "            mid = (l + r) // 2\n",
    "            need = sum(max(0, x - mid) for x in v)\n",
    "            if need <= k1 + k2:\n",
    "                r = mid\n",
    "            else:\n",
    "                l = mid\n",
    "        n = len(v)\n",
    "        t = k1 + k2\n",
    "        for i in range(n):\n",
    "            if v[i] > r:\n",
    "                t -= v[i] - r\n",
    "                v[i] = r\n",
    "        for i in range(n):\n",
    "            if v[i] == r and t:\n",
    "                v[i] -= 1\n",
    "                t -= 1\n",
    "        return sum(max(0, x) ** 2 for x in v)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minSumSquareDiff(self, a: List[int], nums2: List[int], k1: int, k2: int) -> int:\n",
    "        ans, k = 0, k1 + k2\n",
    "        for i in range(len(a)):\n",
    "            a[i] = abs(a[i] - nums2[i])\n",
    "            ans += a[i] * a[i]\n",
    "        if sum(a) <= k:\n",
    "            return 0  # 所有 a[i] 均可为 0\n",
    "        a.sort(reverse=True)\n",
    "        a.append(0)  # 哨兵\n",
    "        for i, v in enumerate(a):\n",
    "            ans -= v * v\n",
    "            j = i + 1\n",
    "            c = j * (v - a[j])\n",
    "            if c < k:\n",
    "                k -= c\n",
    "                continue\n",
    "            v -= k // j\n",
    "            return ans + k % j * (v - 1) * (v - 1) + (j - k % j) * v * v"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minSumSquareDiff(self, a: List[int], nums2: List[int], k1: int, k2: int) -> int:\n",
    "        ans, k = 0, k1 + k2\n",
    "        for i in range(len(a)):\n",
    "            a[i] = abs(a[i] - nums2[i])\n",
    "            ans += a[i] * a[i]\n",
    "        if sum(a) <= k: return 0\n",
    "        a.sort(reverse=True)\n",
    "        a.append(0)\n",
    "        for i, v in enumerate(a):\n",
    "            ans -= v * v\n",
    "            j = i + 1\n",
    "            c = v - a[j]\n",
    "            if k >= j * c:\n",
    "                k -= j * c\n",
    "                continue\n",
    "            v -= k // j\n",
    "            return ans + k % j * (v - 1) * (v - 1) + (j - k % j) * v * v"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minSumSquareDiff(self, nums1: List[int], nums2: List[int], k1: int, k2: int) -> int:\n",
    "        count = [0] * (100001)\n",
    "        for i in range(len(nums1)):\n",
    "            count[abs(nums1[i] - nums2[i])] += 1\n",
    "        k = k1 + k2\n",
    "        for i in range(len(count) - 1, 0, -1):\n",
    "            change = min(k, count[i])\n",
    "            count[i - 1] += change\n",
    "            k -= change\n",
    "            count[i] -= change\n",
    "        res = 0\n",
    "        for i in range(len(count)):\n",
    "            if count[i] != 0:\n",
    "                res += count[i] * i * i \n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minSumSquareDiff(self, nums1: List[int], nums2: List[int], k1: int, k2: int) -> int:\n",
    "        n = len(nums1)\n",
    "        diff = [abs(nums1[i] - nums2[i]) for i in range(n)]\n",
    "        l, r = 0, max(diff)\n",
    "        while l <= r:\n",
    "            m = l + r >> 1\n",
    "            k = k1 + k2\n",
    "            for x in diff:\n",
    "                if x > m:\n",
    "                    k -= x - m\n",
    "                    if k < 0:\n",
    "                        l = m + 1\n",
    "                        break\n",
    "            else:\n",
    "                r = m - 1\n",
    "        if l == 0:\n",
    "            return 0\n",
    "        k = k1 + k2\n",
    "        for x in diff:\n",
    "            if x > l:\n",
    "                k -= x - l\n",
    "        ret = 0\n",
    "        for x in diff:\n",
    "            if x >= l:\n",
    "                if k > 0:\n",
    "                    ret += (l - 1) * (l - 1)\n",
    "                    k -= 1\n",
    "                else:\n",
    "                    ret += l * l\n",
    "            else:\n",
    "                ret += x * x\n",
    "        return ret\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minSumSquareDiff(self, nums1: List[int], nums2: List[int], k1: int, k2: int) -> int:\n",
    "        n=len(nums1)\n",
    "        nums=[abs(nums1[i]-nums2[i]) for i in range(n)]\n",
    "        nums.sort()\n",
    "        res=k1+k2\n",
    "        i=n-1\n",
    "        while i>=0:\n",
    "            if i==0:\n",
    "                minus=(n-i)*nums[i]\n",
    "            else:\n",
    "                minus=(n-i)*(nums[i]-nums[i-1])\n",
    "            if minus <= res:\n",
    "                i-=1\n",
    "                res-=minus\n",
    "            else:\n",
    "                break\n",
    "        if i==-1:\n",
    "            return 0\n",
    "        else:\n",
    "            ans=0\n",
    "            for j in range(0,i):\n",
    "                ans+=nums[j]*nums[j]\n",
    "            ans+=(res%(n-i))*(nums[i]-int(res/(n-i))-1)*(nums[i]-int(res/(n-i))-1)\n",
    "            ans+=(n-i-res%(n-i))*(nums[i]-int(res/(n-i)))*(nums[i]-int(res/(n-i)))\n",
    "            return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minSumSquareDiff(self, a: List[int], nums2: List[int], k1: int, k2: int) -> int:\n",
    "        ans, k = 0, k1 + k2\n",
    "        for i in range(len(a)):\n",
    "            a[i] = abs(a[i] - nums2[i])\n",
    "            ans += a[i] * a[i]\n",
    "        if sum(a) <= k:\n",
    "            return 0  # 所有 a[i] 均可为 0\n",
    "        a.sort(reverse=True)\n",
    "        a.append(0)  # 哨兵\n",
    "        for i, v in enumerate(a):\n",
    "            ans -= v * v\n",
    "            j = i + 1\n",
    "            c = j * (v - a[j])\n",
    "            if c < k:\n",
    "                k -= c\n",
    "                continue\n",
    "            v -= k // j\n",
    "            return ans + k % j * (v - 1) * (v - 1) + (j - k % j) * v * v\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "import heapq\r\n",
    "class Solution:\r\n",
    "    def minSumSquareDiff(self, nums1: List[int], nums2: List[int], k1: int, k2: int) -> int:\r\n",
    "        # count = [0] * (100001)\r\n",
    "        # for i in range(len(nums1)):\r\n",
    "        #     count[abs(nums1[i] - nums2[i])] += 1\r\n",
    "        # k = k1 + k2\r\n",
    "        # for i in range(len(count)-1,0,-1):\r\n",
    "        #     change = min(k,count[i])\r\n",
    "        #     count[i-1] += change\r\n",
    "        #     k -= change\r\n",
    "        #     count[i] -= change\r\n",
    "        # res = 0\r\n",
    "        # for i in range(len(count)):\r\n",
    "        #     if count[i] != 0:\r\n",
    "        #         res += i * i * count[i]\r\n",
    "        # return res\r\n",
    "        ans, k = 0, k1 + k2\r\n",
    "        for i in range(len(nums1)):\r\n",
    "            nums1[i] = abs(nums1[i] - nums2[i])\r\n",
    "            ans += nums1[i] * nums1[i]\r\n",
    "        if sum(nums1) <= k:\r\n",
    "            return 0\r\n",
    "        nums1.sort(reverse=True)\r\n",
    "        nums1.append(0)\r\n",
    "        for i,v in enumerate(nums1):\r\n",
    "            ans -= v * v\r\n",
    "            j = i + 1\r\n",
    "            c = j * (v-nums1[j])\r\n",
    "            if c < k:\r\n",
    "                k -= c\r\n",
    "                continue\r\n",
    "            v -= k // j\r\n",
    "            return ans + k % j * (v-1)**2 + (j-k%j) * v**2\r\n",
    "\r\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minSumSquareDiff(self, nums1: List[int], nums2: List[int], k1: int, k2: int) -> int:\n",
    "        n = len(nums1)\n",
    "        cnt = [0] * 100001\n",
    "        for x, y in zip(nums1, nums2):\n",
    "            print(abs(x-y))\n",
    "            cnt[abs(x-y)]+=1\n",
    "        k = k1 + k2\n",
    "        for i in range(len(cnt)-1, 0, -1):\n",
    "            change = min(k, cnt[i])\n",
    "            cnt[i-1] += change\n",
    "            k -= change\n",
    "            cnt[i] -= change\n",
    "\n",
    "        ans = 0\n",
    "        for i, x in enumerate(cnt):\n",
    "            ans += x * i * i\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "import heapq\r\n",
    "class Solution:\r\n",
    "    def minSumSquareDiff(self, nums1: List[int], nums2: List[int], k1: int, k2: int) -> int:\r\n",
    "        count = [0] * (100001)\r\n",
    "        for i in range(len(nums1)):\r\n",
    "            count[abs(nums1[i] - nums2[i])] += 1\r\n",
    "        k = k1 + k2\r\n",
    "        for i in range(len(count)-1,0,-1):\r\n",
    "            change = min(k,count[i])\r\n",
    "            count[i-1] += change\r\n",
    "            k -= change\r\n",
    "            count[i] -= change\r\n",
    "        res = 0\r\n",
    "        for i in range(len(count)):\r\n",
    "            if count[i] != 0:\r\n",
    "                res += i * i * count[i]\r\n",
    "        return res\r\n",
    "\r\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minSumSquareDiff(self, nums1: List[int], nums2: List[int], k1: int, k2: int) -> int:\n",
    "        cnt = [0] * 100001\n",
    "        for nums1, nums2 in zip(nums1, nums2):\n",
    "            cnt[abs(nums1 - nums2)] += 1\n",
    "        k = k1 + k2\n",
    "        for i in range(len(cnt) - 1, 0, -1):\n",
    "            change = min(k, cnt[i])\n",
    "            cnt[i - 1] += change\n",
    "            k -= change\n",
    "            cnt[i] -= change\n",
    "        res = 0\n",
    "        for i in range(len(cnt)):\n",
    "            if cnt[i] > 0:\n",
    "                res += cnt[i] * i * i\n",
    "        return res\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minSumSquareDiff(self, a: List[int], nums2: List[int], k1: int, k2: int) -> int:\n",
    "        ans, k = 0, k1 + k2\n",
    "        for i in range(len(a)):\n",
    "            a[i] = abs(a[i] - nums2[i])\n",
    "            ans += a[i] * a[i]\n",
    "        if sum(a) <= k:\n",
    "            return 0\n",
    "        a.sort(reverse=True)\n",
    "        a.append(0)\n",
    "        for i, v in enumerate(a):\n",
    "            ans -= v * v\n",
    "            j = i + 1\n",
    "            c = j * (v - a[j])\n",
    "            if c < k:\n",
    "                k -= c\n",
    "                continue\n",
    "            v -= k // j\n",
    "            return ans + k % j * (v - 1) * (v - 1) + (j - k % j) * v * v\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minSumSquareDiff(self, a: List[int], nums2: List[int], k1: int, k2: int) -> int:\n",
    "        ans, k = 0, k1 + k2\n",
    "        for i in range(len(a)):\n",
    "            a[i] = abs(a[i] - nums2[i])\n",
    "            ans += a[i] * a[i]\n",
    "        if sum(a) <= k:\n",
    "            return 0  # 所有 a[i] 均可为 0\n",
    "        a.sort(reverse=True)\n",
    "        a.append(0)  # 哨兵\n",
    "        for i, v in enumerate(a):\n",
    "            ans -= v * v\n",
    "            j = i + 1\n",
    "            if k >= j * (v - a[j]):\n",
    "                k -= j * (v - a[j])\n",
    "                continue\n",
    "            v -= k // j\n",
    "            return ans + k % j * (v - 1) * (v - 1) + (j - k % j) * v * v"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minSumSquareDiff(self, a: List[int], nums2: List[int], k1: int, k2: int) -> int:\n",
    "        ans, k = 0, k1 + k2\n",
    "        for i in range(len(a)):\n",
    "            a[i] = abs(a[i] - nums2[i])\n",
    "            ans += a[i] * a[i]\n",
    "        if sum(a) <= k:\n",
    "            return 0  # 所有 a[i] 均可为 0\n",
    "        a.sort(reverse=True)\n",
    "        a.append(0)  # 哨兵\n",
    "        for i, v in enumerate(a):\n",
    "            ans -= v * v\n",
    "            j = i + 1\n",
    "            c = j * (v - a[j])\n",
    "            if c < k:\n",
    "                k -= c\n",
    "                continue\n",
    "            v -= k // j\n",
    "            return ans + k % j * (v - 1) * (v - 1) + (j - k % j) * v * v\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class BinarySearch:\n",
    "    def __init__(self):\n",
    "        return\n",
    "\n",
    "    @staticmethod\n",
    "    def find_int_left(low: int, high: int, check: Callable) -> int:\n",
    "        # 模板: 整数范围内二分查找，选择最靠左满足check\n",
    "        while low < high - 1:\n",
    "            mid = low + (high - low) // 2\n",
    "            if check(mid):\n",
    "                high = mid\n",
    "            else:\n",
    "                low = mid\n",
    "        return low if check(low) else high\n",
    "\n",
    "    @staticmethod\n",
    "    def find_int_right(low: int, high: int, check: Callable) -> int:\n",
    "        # 模板: 整数范围内二分查找，选择最靠右满足check\n",
    "        while low < high - 1:\n",
    "            mid = low + (high - low) // 2\n",
    "            if check(mid):\n",
    "                low = mid\n",
    "            else:\n",
    "                high = mid\n",
    "        return high if check(high) else low\n",
    "\n",
    "    @staticmethod\n",
    "    def find_float_left(low: float, high: float, check: Callable, error=1e-6) -> float:\n",
    "        # 模板: 浮点数范围内二分查找, 选择最靠左满足check\n",
    "        while low < high - error:\n",
    "            mid = low + (high - low) / 2\n",
    "            if check(mid):\n",
    "                high = mid\n",
    "            else:\n",
    "                low = mid\n",
    "        return low if check(low) else high\n",
    "\n",
    "    @staticmethod\n",
    "    def find_float_right(low: float, high: float, check: Callable, error=1e-6) -> float:\n",
    "        # 模板: 浮点数范围内二分查找, 选择最靠右满足check\n",
    "        while low < high - error:\n",
    "            mid = low + (high - low) / 2\n",
    "            if check(mid):\n",
    "                low = mid\n",
    "            else:\n",
    "                high = mid\n",
    "        return high if check(high) else low\n",
    "\n",
    "\n",
    "class Solution:\n",
    "    def minSumSquareDiff(self, nums1: List[int], nums2: List[int], k1: int, k2: int) -> int:\n",
    "        n = len(nums1)\n",
    "        lst = [abs(nums1[i]-nums2[i]) for i in range(n)]\n",
    "        \n",
    "        def check(x):\n",
    "            return sum(w-x for w in lst if w >= x) <= k1+k2\n",
    "        \n",
    "        k = k1+k2\n",
    "        m = BinarySearch().find_int_left(0, max(lst), check)\n",
    "        for i in range(n):\n",
    "            if lst[i] > m:\n",
    "                k -= lst[i] - m\n",
    "                lst[i] = m\n",
    "        lst.sort(reverse=True)\n",
    "        for i in range(k):\n",
    "            if i >= n:\n",
    "                break\n",
    "            if lst[i]:\n",
    "                lst[i] -= 1\n",
    "        return sum(x*x for x in lst)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from typing import List\n",
    "\n",
    "class Solution:\n",
    "    def minSumSquareDiff(self, a: List[int], nums2: List[int], k1: int, k2: int) -> int:\n",
    "        ans, k = 0, k1 + k2\n",
    "        for i in range(len(a)):\n",
    "            a[i] = abs(a[i] - nums2[i])\n",
    "            ans += a[i] * a[i]\n",
    "        if sum(a) <= k:\n",
    "            return 0  # 所有 a[i] 均可为 0\n",
    "        a.sort(reverse=True)\n",
    "        a.append(0)  # 哨兵\n",
    "        for i, v in enumerate(a):\n",
    "            ans -= v * v\n",
    "            j = i + 1\n",
    "            c = j * (v - a[j])\n",
    "            if c < k:\n",
    "                k -= c\n",
    "                continue\n",
    "            v -= k // j\n",
    "            return ans + k % j * (v - 1) * (v - 1) + (j - k % j) * v * v\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minSumSquareDiff(self, nums1: List[int], nums2: List[int], k1: int, k2: int) -> int:\n",
    "        count = [0] * (100001)\n",
    "        for i in range(len(nums1)):\n",
    "            count[abs(nums1[i] - nums2[i])] += 1\n",
    "        k = k1 + k2\n",
    "        for i in range(len(count) - 1, 0, -1):\n",
    "            change = min(k, count[i])\n",
    "            count[i - 1] += change\n",
    "            k -= change\n",
    "            count[i] -= change\n",
    "        res = 0\n",
    "        for i in range(len(count)):\n",
    "            if count[i] != 0:\n",
    "                res += count[i] * i * i \n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minSumSquareDiff(self, a: List[int], nums2: List[int], k1: int, k2: int) -> int:\n",
    "        ans, k = 0, k1 + k2\n",
    "        for i in range(len(a)):\n",
    "            a[i] = abs(a[i] - nums2[i])\n",
    "            ans += a[i] * a[i]\n",
    "        if sum(a) <= k:\n",
    "            return 0  # 所有 a[i] 均可为 0\n",
    "        a.sort(reverse=True)\n",
    "        a.append(0)  # 哨兵\n",
    "        for i, v in enumerate(a):\n",
    "            ans -= v * v\n",
    "            j = i + 1\n",
    "            c = j * (v - a[j])\n",
    "            if c < k:\n",
    "                k -= c\n",
    "                continue\n",
    "            v -= k // j\n",
    "            return ans + k % j * (v - 1) * (v - 1) + (j - k % j) * v * v\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minSumSquareDiff(self, nums1: List[int], nums2: List[int], k1: int, k2: int) -> int:\n",
    "        arr = [abs(nums1[i] - nums2[i]) for i in range(len(nums1))]\n",
    "        ans = sum(i * i for i in arr)\n",
    "        k = k1 + k2\n",
    "        if k == 0:\n",
    "            return ans\n",
    "        if sum(arr) <= k:\n",
    "            return 0\n",
    "        arr.sort(reverse=True)\n",
    "        arr.append(0)\n",
    "        for i, x in enumerate(arr):\n",
    "            if x == arr[i + 1]:\n",
    "                continue\n",
    "            c = i + 1\n",
    "            ans -= x * x * c\n",
    "            if c * (x - arr[i + 1]) <= k:\n",
    "                ans += c * (arr[i + 1] ** 2)\n",
    "                k -= c * (x - arr[i + 1])\n",
    "                if k <= 0:\n",
    "                    break\n",
    "            else:\n",
    "                d, r = k // c , k % c\n",
    "                ans += r * (x - d - 1) * (x - d - 1) + (c - r) * (x - d) * (x - d)\n",
    "                break\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from typing import List\n",
    "\n",
    "\n",
    "class Solution:\n",
    "    def minSumSquareDiff(self, nums1: List[int], nums2: List[int], k1: int, k2: int) -> int:\n",
    "        n = len(nums1)\n",
    "        nums = []\n",
    "        for i in range(n):\n",
    "            nums.append(abs(nums2[i]-nums1[i]))\n",
    "        # print(nums)\n",
    "        k = k1 + k2\n",
    "        left = 0\n",
    "        right = max(nums)\n",
    "        # print(\"right:{}\".format(right))\n",
    "        while left < right:\n",
    "            mid = (left+right)//2\n",
    "            if self.check(nums,mid,k):\n",
    "               right = mid\n",
    "            else:\n",
    "                left = mid+1\n",
    "        # print(left)\n",
    "        res = 0\n",
    "        for i in range(len(nums)):\n",
    "            if nums[i] > left:\n",
    "                k -= (nums[i]-left)\n",
    "                nums[i] = left\n",
    "        for i in range(len(nums)):\n",
    "            if nums[i]==left and k>0 and nums[i]>0:\n",
    "                # print('i:{},nums:{}'.format(i, nums[i]))\n",
    "                nums[i]-=1\n",
    "                k-=1\n",
    "                # print('i:{},nums:{}'.format(i,nums[i]))\n",
    "        # print(nums)\n",
    "        for num in nums:\n",
    "            res += num**2\n",
    "        return res\n",
    "    def check(self,nums,mid,k):\n",
    "        sum = 0\n",
    "        for num in nums:\n",
    "            if num > mid:\n",
    "                sum += (num-mid)\n",
    "\n",
    "        # print(\"mid:{},sum:{}\".format(mid,sum))\n",
    "        return sum <= k"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minSumSquareDiff(self, nums1: List[int], nums2: List[int], k1: int, k2: int) -> int:\n",
    "        d = [abs(n1-n2) for n1, n2 in zip(nums1, nums2)]\n",
    "        d.sort() \n",
    "        k = k1 + k2 \n",
    "        n = len(d) \n",
    "        def check(x):\n",
    "            tot = 0 \n",
    "            cnt = 0 \n",
    "            for i in range(n-1, -1, -1):\n",
    "                if d[i] >= x:\n",
    "                    tot += d[i] - x \n",
    "                    cnt += 1 \n",
    "                else:\n",
    "                    break \n",
    "            return tot \n",
    "        l, r = 0, max(d)\n",
    "        while l < r:\n",
    "            mid = (l + r) >> 1 \n",
    "            if check(mid) <= k:\n",
    "                r = mid \n",
    "            else:\n",
    "                l = mid + 1\n",
    "        tot = check(l) \n",
    "        r = k - tot \n",
    "        #print(l, tot, r,d) \n",
    "        ans = 0 \n",
    "        for i in range(n-1, -1, -1):\n",
    "            tmp = min(l, d[i]) \n",
    "            if r > 0:\n",
    "                tmp = max(0, tmp-1)\n",
    "                r -= 1\n",
    "            ans += tmp ** 2 \n",
    "        return ans "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minSumSquareDiff(self, nums1: List[int], nums2: List[int], k1: int, k2: int) -> int:\n",
    "        \n",
    "        nums1 = [abs(x - y) for x, y in zip(nums1, nums2)]\n",
    "\n",
    "        k = k1 + k2\n",
    "        total = sum(nums1)\n",
    "        if total <= k:\n",
    "            return 0\n",
    "        \n",
    "        n = max(nums1)\n",
    "        # arr[i]代表差值为i的出现次数\n",
    "        arr = [0] * (n+1)\n",
    "        for num in nums1:\n",
    "            arr[num] += 1\n",
    "        \n",
    "        for i in range(n, -1, -1):\n",
    "            nxt = arr[i-1]\n",
    "            k_minus = arr[i]\n",
    "            if k > k_minus:\n",
    "                k -= k_minus\n",
    "                arr[i-1] += arr[i]\n",
    "                arr[i] = 0\n",
    "            else:\n",
    "                index = i\n",
    "                break\n",
    "        \n",
    "        arr[index] -= k\n",
    "        arr[index-1] += k\n",
    "\n",
    "        return sum(i * i * arr[i] for i in range(n+1))\n",
    "\n",
    "\n",
    "\n",
    "        \n",
    "\n",
    "        \n",
    "        \n",
    "\n",
    "\n",
    "        \n",
    "\n",
    "        \n",
    "\n",
    "\n",
    "        \n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minSumSquareDiff(self, nums1: List[int], nums2: List[int], k1: int, k2: int) -> int:\n",
    "        temp = [abs(nums1[i] - nums2[i]) for i in range(len(nums1))]\n",
    "        k = k1 + k2\n",
    "        if sum(temp) <= k: return 0\n",
    "        \n",
    "        # 二分查找 T 值\n",
    "        l, r, T = 0, max(temp), None\n",
    "        while l <= r:\n",
    "            mid = (l + r) // 2\n",
    "            total = sum(ai - mid for ai in temp if ai >= mid)\n",
    "            if total <= k:\n",
    "                T = mid\n",
    "                r = mid - 1\n",
    "            else:\n",
    "                l = mid + 1\n",
    "\n",
    "        range_sum = lambda x, y: (x + y) * (y - x + 1) // 2\n",
    "        rest = k - sum(ai - T for ai in temp if ai >= T)\n",
    "        ans = 0\n",
    "        for ai in temp:\n",
    "            if ai >= T:\n",
    "                ans += T ** 2\n",
    "            else: \n",
    "                ans += ai ** 2\n",
    "                    \n",
    "        return ans - rest * T ** 2 + rest * (T - 1) ** 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minSumSquareDiff(self, nums1: List[int], nums2: List[int], k1: int, k2: int) -> int:\n",
    "\n",
    "        nums = [abs(x - y) for x, y in zip(nums1, nums2)]\n",
    "\n",
    "        k = k1 + k2\n",
    "        total = sum(nums)\n",
    "        if total <= k:\n",
    "            return 0\n",
    "        \n",
    "        nums.sort(reverse = True)\n",
    "        nums.append(0)\n",
    "\n",
    "\n",
    "        for i, num in enumerate(nums):\n",
    "            nxt = nums[i+1]\n",
    "            k_minus = (i+1) * (num - nxt)\n",
    "            if k_minus < k:\n",
    "                k -= k_minus\n",
    "            else:\n",
    "                index = i\n",
    "                break\n",
    "        \n",
    "        # 此时nums[0]到nums[index]都是nums[index],k还不是0\n",
    "        a, b = k // (index+1), k % (index+1)\n",
    "        # nums[0]到nums[index]每个减a，其中k % index个再减1\n",
    "\n",
    "        ans = (index+1 - b) * (nums[index]-a) * (nums[index]-a)+ b * (nums[index]-a-1) * (nums[index]-a-1)\n",
    "        ans += sum(a*a for a in nums[index+1:])\n",
    "        return ans\n",
    "        \n",
    "\n",
    "        \n",
    "        \n",
    "\n",
    "\n",
    "        \n",
    "\n",
    "        \n",
    "\n",
    "\n",
    "        \n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from math import *\n",
    "from queue import PriorityQueue,Queue\n",
    "from typing import *\n",
    "from heapq import *\n",
    "from functools import lru_cache,cache\n",
    "from string import *\n",
    "from itertools import product\n",
    "from collections import defaultdict\n",
    "# by hangpengjie\n",
    "class Solution:\n",
    "    def minSumSquareDiff(self, nums1: List[int], nums2: List[int], k1: int, k2: int) -> int:\n",
    "        n = len(nums1)\n",
    "        d = []\n",
    "        for x,y in zip(nums1,nums2):\n",
    "            d.append(abs(x - y))\n",
    "        d.sort(reverse=True)\n",
    "        k = k1 + k2\n",
    "        l,r = 0, d[0]\n",
    "        while l < r:\n",
    "            m = (l + r) >> 1\n",
    "            p = k\n",
    "            for x in d:\n",
    "                if x <= m or p < 0:\n",
    "                    break\n",
    "                p = p - (x - m)\n",
    "            if p < 0:\n",
    "                l = m + 1\n",
    "            else:\n",
    "                r = m\n",
    "        if l == 0:\n",
    "            return 0\n",
    "        l -= 1\n",
    "        ans = 0\n",
    "        for i in range(n):\n",
    "            if d[i] > l:\n",
    "                k = k - (d[i] - l)\n",
    "                d[i] = l\n",
    "            else:\n",
    "                break\n",
    "        \n",
    "        for x in d:\n",
    "            if k < 0:\n",
    "                ans += (x + 1) ** 2\n",
    "                k += 1\n",
    "            else:\n",
    "                ans += (x) ** 2\n",
    "            \n",
    "        return ans\n",
    "\n",
    "s = Solution()\n",
    "s.minSumSquareDiff([10,10,10,11,5],[1,0,6,6,1],11,27)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minSumSquareDiff(self, nums1: List[int], nums2: List[int], k1: int, k2: int) -> int:\n",
    "        diff = []\n",
    "        sum_diff = 0\n",
    "        for x, y in zip(nums1, nums2):\n",
    "            if x == y:\n",
    "                continue\n",
    "            res = abs(x - y)\n",
    "            diff.append(res)\n",
    "            sum_diff += res\n",
    "        if sum_diff <= (k1 + k2):\n",
    "            return 0\n",
    "        diff.sort()\n",
    "        sum_diff = sum_diff - (k1 + k2)\n",
    "        n = len(diff)\n",
    "        pre = 0\n",
    "        high = 0\n",
    "        ans = 0\n",
    "        for i, x in enumerate(diff):\n",
    "            if x == high:\n",
    "                ans += high ** 2\n",
    "                continue\n",
    "            res = (x - high) * (n - i)\n",
    "            if res <= sum_diff:\n",
    "                sum_diff -= res\n",
    "                high = x\n",
    "                ans += high ** 2\n",
    "            else:\n",
    "                res1 = sum_diff // (n - i)\n",
    "                res2 = sum_diff % (n - i)\n",
    "                ans += (high + res1) ** 2 * (n - i - res2)\n",
    "                ans += (high + res1 + 1) ** 2 * res2\n",
    "                return ans\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minSumSquareDiff(self, nums1: List[int], nums2: List[int], k1: int, k2: int) -> int:\n",
    "        count = [0] * (100001)\n",
    "        for i in range(len(nums1)):\n",
    "            count[abs(nums1[i] - nums2[i])] += 1\n",
    "        k = k1 + k2\n",
    "        for i in range(len(count) - 1, 0, -1):\n",
    "            change = min(k, count[i])\n",
    "            count[i - 1] += change\n",
    "            k -= change\n",
    "            count[i] -= change\n",
    "        res = 0\n",
    "        for i in range(len(count)):\n",
    "            if count[i] != 0:\n",
    "                res += count[i] * i * i \n",
    "        return res\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minSumSquareDiff(self, nums1: List[int], nums2: List[int], k1: int, k2: int) -> int:\n",
    "        k=k1+k2 \n",
    "        n=len(nums1)\n",
    "        ls=[]\n",
    "        for i in range(n):\n",
    "            ls.append(abs(nums1[i]-nums2[i]))\n",
    "        ls.sort()\n",
    "        l,r=0,ls[-1]\n",
    "        def check(x):\n",
    "            ret=0\n",
    "            for i in range(n):\n",
    "                tp=ls[i]-x\n",
    "                if tp>0:\n",
    "                    ret+=tp\n",
    "            return ret<=k \n",
    "        if check(0):\n",
    "            return 0\n",
    "        while l<r-1:\n",
    "            mid=(l+r)//2\n",
    "            if check(mid):\n",
    "                r=mid\n",
    "            else:\n",
    "                l=mid\n",
    "        if check(l):\n",
    "            last=l\n",
    "        else:\n",
    "            last=r\n",
    "        for i in range(n):\n",
    "            if ls[i]>last:\n",
    "                k-=ls[i]-last\n",
    "                ls[i]=last\n",
    "        ans=0\n",
    "        for i in range(n-1,-1,-1):\n",
    "            if k>0:\n",
    "                ans+=(ls[i]-1)**2\n",
    "                k-=1\n",
    "            else:\n",
    "                ans+=ls[i]**2\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minSumSquareDiff(self, nums1: List[int], nums2: List[int], k1: int, k2: int) -> int:\n",
    "        k1 += k2\n",
    "        for i, (n1, n2) in enumerate(zip(nums1, nums2)):\n",
    "            nums1[i] = abs(n1-n2)\n",
    "        if not k1: \n",
    "            return sum(n*n for n in nums1)\n",
    "        if sum(nums1) <= k1: return 0\n",
    "        nums1.sort()\n",
    "        off = bisect_right(nums1, 0)\n",
    "        if off == len(nums1): return 0\n",
    "        l, r = 0, nums1[-1]\n",
    "        while r - l > 1:\n",
    "            m = (l + r)>>1\n",
    "            delta = sum(n - m for n in nums1[off:] if n > m)\n",
    "            if delta > k1: \n",
    "                l = m\n",
    "            else:\n",
    "                r = m\n",
    "        if r < nums1[-1]:\n",
    "            for i in range(len(nums1)-1, off-1, -1):\n",
    "                n = nums1[i]\n",
    "                if n > r:\n",
    "                    nums1[i] = r\n",
    "                    k1 -= n - r\n",
    "        if k1:\n",
    "            for i in range(1, k1+1):\n",
    "                nums1[-i] -= 1\n",
    "        return sum(n*n  for n in nums1)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minSumSquareDiff(self, nums1: List[int], nums2: List[int], k1: int, k2: int) -> int:\n",
    "        \n",
    "        nums1 = [abs(x - y) for x, y in zip(nums1, nums2)]\n",
    "\n",
    "        k = k1 + k2\n",
    "        total = sum(nums1)\n",
    "        if total <= k:\n",
    "            return 0\n",
    "        \n",
    "        n = max(nums1)\n",
    "        # arr[i]代表差值为i的出现次数\n",
    "        arr = [0] * (n+1)\n",
    "        for num in nums1:\n",
    "            arr[num] += 1\n",
    "        \n",
    "        for i in range(n, -1, -1):\n",
    "            nxt = arr[i-1]\n",
    "            k_minus = arr[i]\n",
    "            if k > k_minus:\n",
    "                k -= k_minus\n",
    "                arr[i-1] += arr[i]\n",
    "                arr[i] = 0\n",
    "            else:\n",
    "                index = i\n",
    "                break\n",
    "        \n",
    "        arr[index] -= k\n",
    "        arr[index-1] += k\n",
    "\n",
    "        return sum(i * i * arr[i] for i in range(n+1))\n",
    "\n",
    "        # 时间复杂度：max(O(max(nums)), O(len(nums1)))\n",
    "\n",
    "\n",
    "\n",
    "        \n",
    "\n",
    "        \n",
    "        \n",
    "\n",
    "\n",
    "        \n",
    "\n",
    "        \n",
    "\n",
    "\n",
    "        \n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "import time\n",
    "\n",
    "class Solution:\n",
    "    def minSumSquareDiff(self, a: List[int], nums2: List[int], k1: int, k2: int) -> int:\n",
    "        ans, k = 0, k1 + k2\n",
    "        for i in range(len(a)):\n",
    "            a[i] = abs(a[i] - nums2[i])\n",
    "            ans += a[i] * a[i]\n",
    "        if sum(a) <= k:\n",
    "            return 0\n",
    "        a.sort(reverse=True)\n",
    "        a.append(0)\n",
    "        for i, v in enumerate(a):\n",
    "            ans -= v * v\n",
    "            j = i + 1\n",
    "            c = j * (v - a[j])\n",
    "            if c < k:\n",
    "                k -= c\n",
    "                continue\n",
    "            v -= k // j\n",
    "            time.sleep(0.3)\n",
    "            return ans + k % j * (v - 1) * (v - 1) + (j - k % j) * v * v\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from math import *\n",
    "from queue import PriorityQueue,Queue\n",
    "from typing import *\n",
    "from heapq import *\n",
    "from functools import lru_cache,cache\n",
    "from string import *\n",
    "from itertools import product\n",
    "from collections import defaultdict\n",
    "# by hangpengjie\n",
    "class Solution:\n",
    "    def minSumSquareDiff(self, nums1: List[int], nums2: List[int], k1: int, k2: int) -> int:\n",
    "        n = len(nums1)\n",
    "        d = []\n",
    "        for x,y in zip(nums1,nums2):\n",
    "            d.append(abs(x - y))\n",
    "        d.sort(reverse=True)\n",
    "        k = k1 + k2\n",
    "        l,r = 0, d[0]\n",
    "        # 二分确定数组d中能达到的最大值\n",
    "        while l < r:\n",
    "            m = (l + r) >> 1\n",
    "            p = k\n",
    "            for x in d:\n",
    "                if x <= m or p < 0:\n",
    "                    break\n",
    "                p = p - (x - m)\n",
    "            if p < 0:\n",
    "                l = m + 1\n",
    "            else:\n",
    "                r = m\n",
    "        # 可以到达0直接返回\n",
    "        if l == 0:\n",
    "            return 0\n",
    "        # 整体不可以达到 l - 1， 最大值一定是 l ，此时k可能不为0，部分为 l 的值可能可以达到 l - 1\n",
    "        l -= 1\n",
    "        ans = 0\n",
    "        for i in range(n):\n",
    "            if d[i] > l:\n",
    "                k = k - (d[i] - l)\n",
    "                d[i] = l\n",
    "            else:\n",
    "                break\n",
    "        \n",
    "        for x in d:\n",
    "            if k < 0:\n",
    "                ans += (x + 1) ** 2\n",
    "                k += 1\n",
    "            else:\n",
    "                ans += (x) ** 2\n",
    "            \n",
    "        return ans\n",
    "\n",
    "s = Solution()\n",
    "s.minSumSquareDiff([10,10,10,11,5],[1,0,6,6,1],11,27)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minSumSquareDiff(self, nums1: List[int], nums2: List[int], k1: int, k2: int) -> int:\n",
    "        k=k1+k2\n",
    "        n=len(nums1)\n",
    "        dif=[0]*n\n",
    "        for i in range(n):\n",
    "            dif[i]=abs(nums1[i]-nums2[i])\n",
    "        dif.sort(reverse=True)\n",
    "        def check(t):\n",
    "            ans=0\n",
    "            for i in range(n):\n",
    "                if dif[i]>t:\n",
    "                    ans+=dif[i]-t\n",
    "                else:\n",
    "                    break\n",
    "            if ans<=k:\n",
    "                return True\n",
    "            else:\n",
    "                return False\n",
    "        l=0\n",
    "        r=max(dif)\n",
    "        #print(dif)\n",
    "        while l+1<r:\n",
    "            mid=(l+r)//2\n",
    "            if check(mid):\n",
    "                r=mid\n",
    "            else:\n",
    "                l=mid\n",
    "        \n",
    "        if check(l):\n",
    "            y=l\n",
    "        else:\n",
    "            y=r\n",
    "        print(dif)\n",
    "        print(y)\n",
    "        ans=0\n",
    "        if y==0:\n",
    "            return 0\n",
    "\n",
    "        for i in range(n):\n",
    "            if dif[i]>y:\n",
    "                ans+=dif[i]-y\n",
    "                dif[i]=y\n",
    "        k =k-ans\n",
    "        print(k)\n",
    "        dif.sort(reverse=True)\n",
    "        for i in range(k):\n",
    "            dif[i]=dif[i]-1\n",
    "        res=0\n",
    "        for i in range(n):\n",
    "            res+=dif[i]**2\n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minSumSquareDiff(self, nums1: List[int], nums2: List[int], k1: int, k2: int) -> int:\n",
    "        n = len(nums1)\n",
    "        diff = [abs(nums1[i] - nums2[i]) for i in range(n)]\n",
    "        print(diff)\n",
    "        # diff.sort()\n",
    "        l, r = 0, max(diff)\n",
    "        while l <= r:\n",
    "            m = l + r >> 1\n",
    "            k = k1 + k2\n",
    "            for x in diff:\n",
    "                if x > m:\n",
    "                    k -= x - m\n",
    "                    if k < 0:\n",
    "                        l = m + 1\n",
    "                        break\n",
    "            else:\n",
    "                r = m - 1\n",
    "        if l == 0:\n",
    "            return 0\n",
    "        k = k1 + k2\n",
    "        for x in diff:\n",
    "            if x > l:\n",
    "                k -= x - l\n",
    "        # print(k, k1 + k2, diff)\n",
    "        ret = 0\n",
    "        for x in diff:\n",
    "            if x >= l:\n",
    "                if k > 0:\n",
    "                    ret += (l - 1) * (l - 1)\n",
    "                    k -= 1\n",
    "                else:\n",
    "                    ret += l * l\n",
    "            else:\n",
    "                ret += x * x\n",
    "        return ret\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minSumSquareDiff(self, nums1: List[int], nums2: List[int], k1: int, k2: int) -> int:\n",
    "        diff = []\n",
    "        for i in range(len(nums1)):\n",
    "            diff.append(abs(nums1[i] - nums2[i]))\n",
    "        times = k1 + k2\n",
    "        diff.sort(reverse=True)\n",
    "        count_map = collections.defaultdict()\n",
    "        count_map[diff[0]] = 1\n",
    "        for i in range(1, len(diff)):\n",
    "            cur_diff = abs(diff[i] - diff[i-1])\n",
    "            last_count = count_map[diff[i-1]]\n",
    "            need_times = cur_diff * last_count\n",
    "            if times >= need_times:\n",
    "                times -= need_times\n",
    "                count_map[diff[i]] = count_map[diff[i-1]] + 1\n",
    "            else:\n",
    "                avg_sub = times // count_map[diff[i-1]]\n",
    "                rest = times % count_map[diff[i-1]]\n",
    "                ans = 0\n",
    "                for j in range(i, len(diff)):\n",
    "                    ans += diff[j] * diff[j]\n",
    "                ans += (count_map[diff[i-1]] - rest) * pow(diff[i-1] - avg_sub, 2) + rest * pow(diff[i-1] - avg_sub - 1, 2) \n",
    "                return ans \n",
    "        count = count_map[diff[i]]\n",
    "        avg_sub = times // count\n",
    "        rest = times % count\n",
    "        if diff[i] * count > times:\n",
    "            return pow(diff[i] - avg_sub, 2) * (count - rest) + rest * pow(diff[i] - avg_sub - 1, 2)\n",
    "        return 0\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "        # heap = []\n",
    "        # for i in range(len(nums1)):\n",
    "        #     heap.append(-abs(nums1[i] - nums2[i]))\n",
    "        # times = k1 + k2\n",
    "        # heapq.heapify(heap)\n",
    "        # for i in range(times):\n",
    "        #     pop_val = heapq.heappop(heap)\n",
    "        #     if pop_val == 0:\n",
    "        #         return 0\n",
    "        #     heapq.heappush(heap, pop_val + 1)\n",
    "        # ans = 0\n",
    "        # for num in heap:\n",
    "        #     ans += num * num\n",
    "        # return ans\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from typing import List\n",
    "\n",
    "\n",
    "class Solution:\n",
    "    def minSumSquareDiff(self, nums1: List[int], nums2: List[int], k1: int, k2: int) -> int:\n",
    "        n = len(nums1)\n",
    "        nums = []\n",
    "        for i in range(n):\n",
    "            nums.append(abs(nums2[i]-nums1[i]))\n",
    "        print(nums)\n",
    "        k = k1 + k2\n",
    "        left = 0\n",
    "        right = max(nums)\n",
    "        print(\"right:{}\".format(right))\n",
    "        while left < right:\n",
    "            mid = (left+right)//2\n",
    "            if self.check(nums,mid,k):\n",
    "               right = mid\n",
    "            else:\n",
    "                left = mid+1\n",
    "        print(left)\n",
    "        res = 0\n",
    "        for i in range(len(nums)):\n",
    "            if nums[i] > left:\n",
    "                k -= (nums[i]-left)\n",
    "                nums[i] = left\n",
    "        for i in range(len(nums)):\n",
    "            if nums[i]==left and k>0 and nums[i]>0:\n",
    "                print('i:{},nums:{}'.format(i, nums[i]))\n",
    "                nums[i]-=1\n",
    "                k-=1\n",
    "                print('i:{},nums:{}'.format(i,nums[i]))\n",
    "        print(nums)\n",
    "        for num in nums:\n",
    "            res += num**2\n",
    "        return res\n",
    "    def check(self,nums,mid,k):\n",
    "        sum = 0\n",
    "        for num in nums:\n",
    "            if num > mid:\n",
    "                sum += (num-mid)\n",
    "\n",
    "        print(\"mid:{},sum:{}\".format(mid,sum))\n",
    "        return sum <= k"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minSumSquareDiff(self, nums1: List[int], nums2: List[int], k1: int, k2: int) -> int:\n",
    "        # 首先可以确定的是，肯定是集中修改某一部分大的值\n",
    "        k = k1 + k2\n",
    "        nums = sorted([0] + [abs(x-y) for x, y in zip(nums1, nums2)], reverse=True)\n",
    "        \n",
    "        s = 0\n",
    "        for i, x in enumerate(nums):\n",
    "            if s - i * x >= k:\n",
    "                break\n",
    "            s += x\n",
    "            \n",
    "        print(nums, k)\n",
    "        if i == 0:\n",
    "            return sum(x**2 for x in nums)\n",
    "        if s - i * x < k:\n",
    "            return 0\n",
    "\n",
    "        # 这个时候开始考虑比较通用的方法\n",
    "        v = nums[i-1]\n",
    "        for j in range(i):\n",
    "            k -= nums[j] - v\n",
    "            nums[j] = v\n",
    "        v1 = k // i\n",
    "        k = k % i\n",
    "        \n",
    "        for j in range(i):\n",
    "            if j < k:\n",
    "                nums[j] -= v1 +1\n",
    "            else:\n",
    "                nums[j] -= v1\n",
    "        \n",
    "        return sum(x**2 for x in nums)\n",
    "        \n",
    "        \n",
    "        \n",
    "        \n",
    "\n",
    "            \n",
    "                \n",
    "                \n",
    "                \n",
    "            \n",
    "            \n",
    "        \n",
    "                        \n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minSumSquareDiff(self, nums1: List[int], nums2: List[int], k1: int, k2: int) -> int:\n",
    "        n = len(nums1)\n",
    "        state = [0] * n\n",
    "        for i in range(0,n):\n",
    "            state[i] = abs(nums1[i] - nums2[i])\n",
    "        k = k1 + k2\n",
    "        state.sort(reverse=True)\n",
    "        state += [0]\n",
    "        m, mod = 0, 0\n",
    "        temp_max = state[0]\n",
    "        index = 0\n",
    "        while index < n and k > 0:\n",
    "            if k >= (state[index] - state[index + 1]) * (index + 1):\n",
    "                temp_max = state[index + 1]\n",
    "                k -= (state[index] - state[index + 1]) * (index + 1)\n",
    "                index += 1\n",
    "            else:\n",
    "                m = k // (index + 1)\n",
    "                mod = k % (index + 1)\n",
    "                k = 0\n",
    "        new_state = [temp_max - m] * (index + 1 - mod) + [temp_max - m - 1] * (mod) + state[index+1:]\n",
    "        res = 0\n",
    "        for s in new_state:\n",
    "            res += s**2\n",
    "        return res\n",
    "\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minSumSquareDiff(self, nums1: List[int], nums2: List[int], k1: int, k2: int) -> int:\n",
    "        L = [a-b for a,b in zip(nums1,nums2)]\n",
    "        n1 = [-a for a in L if a < 0]\n",
    "        n2 = [a for a in L if a > 0]\n",
    "        n1.extend(n2)\n",
    "        n1.sort()\n",
    "        def f(L,kk):\n",
    "            l = len(L)\n",
    "            if sum(L) <= kk:return 0\n",
    "            def check(k):\n",
    "                cnt = 0;s = 0\n",
    "                for i in L:\n",
    "                    if i >= k:\n",
    "                        s += i-k;cnt += 1\n",
    "                # print(k,s+cnt,kk)\n",
    "                if s + cnt <= kk:return True\n",
    "                return False\n",
    "            shold = bisect_left(range(10**5+1),True,key = check)-1\n",
    "            # print(L,shold)\n",
    "            res = 0;cnt = 0;s = 0\n",
    "            for i in L:\n",
    "                if i < shold:res += i**2;continue\n",
    "                cnt += 1;s += i-shold\n",
    "            res += (kk - s) * (shold-1) ** 2\n",
    "            res += (cnt - kk + s) * shold ** 2\n",
    "            return res\n",
    "        # print(k1,k2,n1,n2)\n",
    "        return f(n1,k1+k2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minSumSquareDiff(self, nums1: List[int], nums2: List[int], k1: int, k2: int) -> int:\n",
    "        def lt233(nums1,nums2,k1,k2):\n",
    "            n = len(nums1)\n",
    "            ll = [abs(nums1[i] - abs(nums2[i])) for i in range(n)]\n",
    "            if sum(ll)<=(k1+k2):return 0\n",
    "            ll.sort(reverse=True)\n",
    "            dif=[ll[i]-ll[i+1] for i in range(n-1)]\n",
    "            # dif.append(ll[-1])\n",
    "            k=k1+k2\n",
    "            if dif[0]>=k:\n",
    "                ll[0]-=k\n",
    "                return sum([x**2 for x in ll])\n",
    "            sum_,ct=0,0\n",
    "            for i in range(n-1):\n",
    "                ct+=1\n",
    "                sum_+=ct*dif[i]\n",
    "                if sum_>=k:break\n",
    "            if sum_>k:\n",
    "                flag=ll[ct-1]\n",
    "                remain=k-(sum_-ct*dif[i])\n",
    "                h2=flag-remain//ct\n",
    "                ct2=ct-remain%ct\n",
    "                h1=h2-1\n",
    "                ct1=ct-ct2\n",
    "                return h1**2*ct1+h2**2*ct2+sum([x**2 for x in ll[ct:]])\n",
    "            if sum_==k:\n",
    "                flag=ll[ct]\n",
    "                return flag**2*ct+sum([x**2 for x in ll[ct:]])\n",
    "            if sum_<k:\n",
    "                flag=ll[-1]\n",
    "                remain=k-sum_\n",
    "                h1=flag-remain//n\n",
    "                h2=h1-1\n",
    "                ct2=remain%n\n",
    "                h1=h2+1\n",
    "                ct1=n-ct2\n",
    "                return h1 ** 2 * ct1 + h2 ** 2 * ct2\n",
    "        return lt233(nums1,nums2,k1,k2)\n",
    "        \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minSumSquareDiff(self, nums1: List[int], nums2: List[int], k1: int, k2: int) -> int:\n",
    "        l = len(nums1)\n",
    "        diff = sorted([abs(nums1[i] - nums2[i]) for i in range(l) if nums1[i] != nums2[i]])\n",
    "        k = k1 + k2\n",
    "        s = sum(diff)#差值和\n",
    "        if s <= k:\n",
    "            return 0\n",
    "        for i, n in enumerate(diff):\n",
    "            if s - (len(diff) - i) * n <= k:#判断以n为最大值是否可行，只要剩余所有元素改为n的代价不大于k即可\n",
    "                k -= s - (len(diff) - i) * n#更新k\n",
    "                m = k // (len(diff) - i)#最大值可以减少的值\n",
    "                k %= (len(diff) - i)#有多少最大值可以额外减1\n",
    "                break\n",
    "            s -= n#更新差值和\n",
    "        diff = diff[:i] + [n - m - 1] * k + [n - m] * (len(diff) - i - k) #由于这里只修改最大值，所以i之前的不变，i之后的都为n - m, 且有k个可以额外减1\n",
    "        return sum([i ** 2 for i in diff])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minSumSquareDiff(self, nums1: List[int], nums2: List[int], k1: int, k2: int) -> int:\n",
    "        l = len(nums1)\n",
    "        diff = sorted([abs(nums1[i] - nums2[i]) for i in range(l) if nums1[i] != nums2[i]])\n",
    "        k = k1 + k2\n",
    "        s = sum(diff)#差值和\n",
    "        if s <= k:\n",
    "            return 0\n",
    "        for i, n in enumerate(diff):\n",
    "            if s - (len(diff) - i) * n <= k:#判断以n为最大值是否可行，只要剩余所有元素改为n的代价不大于k即可\n",
    "                k -= s - (len(diff) - i) * n#更新k\n",
    "                m = k // (len(diff) - i)#最大值可以减少的值\n",
    "                k %= (len(diff) - i)#有多少最大值可以额外减1\n",
    "                break\n",
    "            s -= n#更新差值和\n",
    "        diff = diff[:i] + [n - m - 1] * k + [n - m] * (len(diff) - i - k) #由于这里只修改最大值，所以i之前的不变，i之后的都为n - m, 且有k个可以额外减1\n",
    "        return sum([i ** 2 for i in diff])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minSumSquareDiff(self, nums1: List[int], nums2: List[int], k1: int, k2: int) -> int:\n",
    "        substracts = [nums1[i] - nums2[i] if nums1[i] > nums2[i] else nums2[i] - nums1[i] for i in range(len(nums1))]\n",
    "        substracts.sort(reverse=True)\n",
    "        substracts.append(0)\n",
    "\n",
    "        k = k1 + k2\n",
    "        \n",
    "        if sum(substracts) <= k:\n",
    "            return 0\n",
    "\n",
    "        c = 0\n",
    "        for i in range(1, len(substracts) + 1):\n",
    "            c += i * (substracts[i-1] - substracts[i])\n",
    "            if c >= k:\n",
    "                last_index = i\n",
    "                last_sum = c - i * (substracts[i-1] - substracts[i])\n",
    "                break\n",
    "\n",
    "        for i in range(last_index):\n",
    "            if i < (k - last_sum) % last_index:\n",
    "                substracts[i] = substracts[last_index - 1] - (k - last_sum) // last_index - 1\n",
    "            else:\n",
    "                substracts[i] = substracts[last_index - 1] - (k - last_sum) // last_index\n",
    "\n",
    "        return sum([x * x for x in substracts])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minSumSquareDiff(self, nums1: List[int], nums2: List[int], k1: int, k2: int) -> int:\n",
    "        n = len(nums1)\n",
    "        k = k1 + k2\n",
    "        arr = sorted([abs(x - y) for x, y in zip(nums1, nums2)], key=lambda z: -z)\n",
    "        arr.append(0)\n",
    "        #print(f'arr:{arr}')\n",
    "        psm = 0\n",
    "        for i, x in enumerate(arr):\n",
    "            sm = psm + (arr[i - 1] - x) * i if i > 0 else 0\n",
    "            #print(f'sm:{sm},i:{i},x:{x}')\n",
    "            if sm > k:\n",
    "                t = k - psm\n",
    "                a = t // i\n",
    "                c = t % i\n",
    "                #print(f'a:{a},c:{c},t:{t}')\n",
    "                ans = 0\n",
    "                for j in range(c):\n",
    "                    ans += (arr[i - 1] - a - 1) ** 2\n",
    "                #print(f'ans:{ans}')\n",
    "                for j in range(c, i):\n",
    "                    ans += (arr[i - 1] - a) ** 2\n",
    "                #print(f'ans:{ans}')\n",
    "                for j in range(i, n):\n",
    "                    ans += arr[j] ** 2\n",
    "                #print(f'ans:{ans}')\n",
    "                return ans\n",
    "            psm = sm\n",
    "        return 0\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minSumSquareDiff(self, nums1: List[int], nums2: List[int], k1: int, k2: int) -> int:\n",
    "        n = len(nums1)\n",
    "        nums = sorted([abs(a - b) for a, b in zip(nums1, nums2)], key=lambda x: -x)\n",
    "        nums.append(0)\n",
    "        k = k1 + k2\n",
    "        for i, num in enumerate(nums):\n",
    "            t = 0 if i == 0 else i * (nums[i - 1] - num)\n",
    "            if t <= k:\n",
    "                k -= t\n",
    "            else:\n",
    "                a = k // i \n",
    "                b = k % i\n",
    "                return (nums[i - 1] - a - 1) ** 2 * b + (nums[i - 1] - a) ** 2 * (i - b) + sum(nums[j] ** 2 for j in range(i,n))\n",
    "        return 0\n",
    "                "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minSumSquareDiff(self, nums1: List[int], nums2: List[int], k1: int, k2: int) -> int:\n",
    "        n, k = len(nums1), k1 + k2\n",
    "        nums = sorted([abs(a - b) for a, b in zip(nums1, nums2)], key=lambda x: -x) + [0]\n",
    "        for i, num in enumerate(nums):\n",
    "            t = 0 if i == 0 else i * (nums[i - 1] - num)\n",
    "            if t <= k:\n",
    "                k -= t\n",
    "                continue\n",
    "            a, b = k // i, k % i\n",
    "            return (nums[i - 1] - a - 1) ** 2 * b + (nums[i - 1] - a) ** 2 * (i - b) + sum(nums[j] ** 2 for j in range(i,n))\n",
    "        return 0\n",
    "                "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minSumSquareDiff(self, nums1: List[int], nums2: List[int], k1: int, k2: int) -> int:\n",
    "        n, k = len(nums1), k1 + k2\n",
    "        nums = sorted([abs(a - b) for a, b in zip(nums1, nums2)], key=lambda x: -x) + [0]\n",
    "        for i, num in enumerate(nums):\n",
    "            t = 0 if i == 0 else i * (nums[i - 1] - num)\n",
    "            if t <= k:\n",
    "                k -= t\n",
    "            else:\n",
    "                a, b = k // i, k % i\n",
    "                return (nums[i - 1] - a - 1) ** 2 * b + (nums[i - 1] - a) ** 2 * (i - b) + sum(nums[j] ** 2 for j in range(i,n))\n",
    "        return 0\n",
    "                "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minSumSquareDiff(self, nums1: List[int], nums2: List[int], k1: int, k2: int) -> int:\n",
    "        n, k = len(nums1), k1 + k2\n",
    "        nums = sorted([abs(a - b) for a, b in zip(nums1, nums2)], key=lambda x: -x) + [0]\n",
    "        for i, num in enumerate(nums):\n",
    "            t = 0 if i == 0 else i * (nums[i - 1] - num)\n",
    "            if t <= k:\n",
    "                k -= t\n",
    "                continue\n",
    "            a, b = k // i, k % i\n",
    "            ans = (nums[i - 1] - a - 1) ** 2 * b \n",
    "            ans += (nums[i - 1] - a) ** 2 * (i - b) \n",
    "            ans += sum(nums[j] ** 2 for j in range(i,n))\n",
    "            return ans\n",
    "        return 0\n",
    "                "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\r\n",
    "    def minSumSquareDiff(self, nums1: List[int], nums2: List[int], k1: int, k2: int) -> int:\r\n",
    "        n = len(nums1) \r\n",
    "        d = [abs(nums1[i] - nums2[i]) for i in range(n)]\r\n",
    "        if sum(d) <= k1 + k2 :\r\n",
    "            return 0\r\n",
    "        d.append(0)\r\n",
    "        d.sort(reverse=True) \r\n",
    "        l, r = 0, n \r\n",
    "        ss = list(accumulate(d, initial=0))\r\n",
    "        def check(x): \r\n",
    "            return  k1 + k2 >= ss[x] - (x * d[x])\r\n",
    "        while l + 1 < r:\r\n",
    "            mid = (l + r) // 2 \r\n",
    "            if check(mid):\r\n",
    "                l = mid\r\n",
    "            else:\r\n",
    "                r = mid \r\n",
    "        left = (k1 + k2) - (ss[l] - (l * d[l]))\r\n",
    "        # print(l, left, d)\r\n",
    "        less = left // (l + 1)\r\n",
    "        cnt = left % (l + 1)\r\n",
    "        res = cnt * pow(d[l] - less - 1, 2) + (l - cnt + 1) * pow(d[l] - less, 2) + sum(d[x] ** 2 for x in range(l + 1, n)) \r\n",
    "        return res \r\n",
    "\r\n",
    "\r\n",
    "\r\n",
    "            \r\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minSumSquareDiff(self, nums1: List[int], nums2: List[int], k1: int, k2: int) -> int:\n",
    "        n = len(nums1)\n",
    "        diff = [abs(nums1[i] - nums2[i]) for i in range(n)]\n",
    "        t = k1 + k2\n",
    "        if t >= sum(diff):\n",
    "            return 0\n",
    "        cnt = Counter(diff)\n",
    "        l = sorted(cnt.keys(), reverse=True)\n",
    "        if l[-1] != 0:\n",
    "            l.append(0)\n",
    "        # print(l)\n",
    "        for i in range(len(l) - 1):\n",
    "            x, y = l[i], l[i + 1]\n",
    "            if (cur := (x - y) * cnt[x]) <= t:\n",
    "                t -= cur\n",
    "                cnt[y] += cnt[x]\n",
    "                del cnt[x]\n",
    "            else:\n",
    "                a, b = divmod(t, cnt[x])\n",
    "                cx = cnt[x]\n",
    "                del cnt[x]\n",
    "                cnt[x - a] += cx - b\n",
    "                cnt[x - a - 1] += b\n",
    "                break\n",
    "        return sum(k ** 2 * v for k, v in cnt.items())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minSumSquareDiff(self, nums1: List[int], nums2: List[int], k1: int, k2: int) -> int:\n",
    "        n = len(nums1)\n",
    "        k = k1 + k2\n",
    "        arr = sorted([abs(x - y) for x, y in zip(nums1, nums2)], key=lambda z: -z)\n",
    "        arr.append(0)\n",
    "        psm = 0\n",
    "        for i, x in enumerate(arr):\n",
    "            sm = psm + (arr[i - 1] - x) * i if i > 0 else 0\n",
    "            if sm > k:\n",
    "                t = k - psm\n",
    "                a = t // i\n",
    "                c = t % i\n",
    "                ans = 0\n",
    "                ans += (arr[i - 1] - a - 1) ** 2 * c\n",
    "                ans += (arr[i - 1] - a) ** 2 * (i - c)\n",
    "                ans += sum(arr[j] ** 2 for j in range(i, n))\n",
    "                return ans\n",
    "            psm = sm\n",
    "        return 0\n",
    "\n",
    "\n"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 2
}
